package top.huhuiyu.springboot3.service.impl;

import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import top.huhuiyu.springboot3.aop.BucketLimit;
import top.huhuiyu.springboot3.aop.BucketLimitMode;
import top.huhuiyu.springboot3.aop.Limit;
import top.huhuiyu.springboot3.service.LimitService;
import top.huhuiyu.springboot3.service.RedisService;
import top.huhuiyu.springboot3.util.IpUtils;

import java.util.Optional;

@Service
@Slf4j
@RequiredArgsConstructor
public class LimitServiceImpl implements LimitService {
    private final RedisService redisService;

    @Override
    public boolean checkLimit(Limit limit) {
        // 组织校验的key
        String key = limit.key();
        if ("ip".equals(limit.mode())) {
            key = key + "_" + IpUtils.getIpAddress();
        }
        // 存在数据表示被限流
        if (redisService.has(key)) {
            return false;
        }
        // 不存在就保存数据
        redisService.save(key, "", limit.timeout());
        return true;
    }

    @Override
    public boolean checkLimit(BucketLimit limit) {
        String key = limit.key();
        if (BucketLimitMode.IP.equals(limit.mode())) {
            key += "_" + IpUtils.getIpAddress();
        }
        Integer count = redisService.get(key, Integer.class);
        count = Optional.ofNullable(count).orElse(0);
        // 桶限流模式需要判断次数
        if (!redisService.has(key)) {
            redisService.save(key, count, limit.timeout());
        }
        if (count < limit.capacity()) {
            // 增加次数后更新回去
            count++;
            redisService.update(key, count);
            return true;
        }
        return false;
    }
}
